from khash import newID
from const import TOKEN_UPDATE
from BitTornado.clock import clock

class TokensHandler:
    """
    Handles generating new tokens, checking tokens etc...
    """
    def __init__(self, factory):
        self.factory = factory
        self.tokens = {}
        self.secret = None
        self.newSecret()      
        self.factory.rawserver.add_task(self.cleanTokens, TOKEN_UPDATE)

    def newSecret(self):
        """
        Generate new secret every 5 minutes
        """
        self.oldSecret = self.secret
        self.secret = newID(digest=False)
        self.factory.rawserver.add_task(self.newSecret, TOKEN_UPDATE)

    def cleanTokens(self):
        """
        Remove old tokens
        """
        for id in self.tokens.keys():
            try:
                if clock() - self.tokens[id][1] > TOKEN_UPDATE:
                    del self.tokens[id]
            except KeyError:
                pass
        self.factory.rawserver.add_task(self.cleanTokens, TOKEN_UPDATE)

    def tokenToSend(self, infohash, ip, port, secret = None):
        """
        Returns the token that should be sent to node with that ip and port
        """
        if not secret:
            secret = self.secret
        token = secret.copy()
        token.update(''.join([ip, str(port), infohash]))
        return token.digest()

    def checkToken(self, token, infohash, ip, port):
        """
        Checks if the token we got match to the token we sent
        """
        if token == self.tokenToSend(infohash, ip, port):
            return True
        if self.oldSecret and token == self.tokenToSend(infohash, ip, port, self.oldSecret):
            return True
        return False

    def gotToken(self, id, token):
        """
        Stores tokens
        """
        self.tokens[id] = (token, clock())

    def getToken(self, id):
        """
        Returns the token stored for that id
        """
        if id in self.tokens:
            return self.tokens[id][0]
        return None
